#include "topology_sort.h"
#include <algorithm>
#include <queue>


namespace vkinfer
{
    struct GraphNode
    {
        std::string name;
        std::vector<GraphNode*> prev;
        std::vector<GraphNode*> next;
        GraphNode(const std::string& name) : name(name) {}

        bool has_next(GraphNode* p)
        {
            if (std::find(next.begin(), next.end(), p) == next.end())
                return false;
            return true;
        }

        bool has_prev(GraphNode* p)
        {
            if (std::find(prev.begin(), prev.end(), p) == prev.end())
                return false;
            return true;
        }
    };

    void connect_nodes(GraphNode* p1, GraphNode* p2)
    {
        if (!p1->has_next(p2)) p1->next.push_back(p2);
        if (!p2->has_prev(p1)) p2->prev.push_back(p1);
    }

    struct GraphEdge
    {
        std::string name;
        std::string from;
        std::vector<std::string> to;
        GraphEdge(const std::string& n) : name(n) {}
    };

    template<typename T>
    bool insert_if_not_exists(std::vector<T>& collection, const T& data)
    {
        // return true if data is inserted.
        auto it = std::find(collection.begin(), collection.end(), data);
        if (it == collection.end())
        {
            collection.push_back(data);
            return true;
        }
        return false;
    }

    bool contains_any(const std::vector<std::string>& data, const std::vector<std::string>& samples)
    {
        for (auto& sample_name : samples)
            if (std::find(data.begin(), data.end(), sample_name) != data.end())
                return true;
        return false;
    }

    bool contains_all(const std::vector<std::string>& data, const std::vector<std::string>& samples)
    {
        for (auto& sample_name : samples)
            if (std::find(data.begin(), data.end(), sample_name) == data.end())
                return false;
        return true;
    }

    std::vector<std::string> topology_sort(
        const std::map<std::string, std::shared_ptr<OnnxNodeInfo>> node_infos,
        std::vector<std::string> input_names)
    {
        // construct graph
        std::map<std::string, GraphNode*> nodes;
        std::map<std::string, GraphEdge*> edges;
        // std::vector<GraphEdge*> edges;
        for (auto& it : node_infos)
        {
            std::shared_ptr<OnnxNodeInfo> node_info = it.second;
            std::string node_name = node_info->name;
            auto& input_names = node_info->input_names;
            auto& output_names = node_info->output_names;
            GraphNode* node = new GraphNode(node_name);
            nodes.insert(std::make_pair(node_name, node));
            // create edges
            for (auto& name : input_names)
            {
                if (edges.find(name) == edges.end())
                {
                    GraphEdge* edge = new GraphEdge(name);
                    edges.insert(std::make_pair(name, edge));
                }
                edges[name]->to.push_back(node_name);
            }
            for (auto& name : output_names)
            {
                if (edges.find(name) == edges.end())
                {
                    GraphEdge* edge = new GraphEdge(name);
                    edges.insert(std::make_pair(name, edge));
                }
                edges[name]->from = node_name;
            }
        }

        // connect nodes
        for (auto& it : edges)
        {
            std::string from = it.second->from;
            if (from.size() == 0) continue;
            std::vector<std::string>& to = it.second->to;
            GraphNode* p1 = nodes[from];
            for (auto& to_node : to)
            {
                GraphNode* p2 = nodes[to_node];
                connect_nodes(p1, p2);
            }
        }

        std::vector<std::string> sorted_nodes;
        std::queue<std::string> in_stock;
        // get input nodes
        for (auto& it : node_infos)
        {
            if (contains_all(input_names, it.second->input_names))
            {
                sorted_nodes.push_back(it.first);
                in_stock.push(it.first);
            }
        }
        // scan graph
        while (!in_stock.empty())
        {
            std::string name = in_stock.front();
            in_stock.pop();
            // scan its child
            std::vector<GraphNode*>& childs = nodes[name]->next;
            for (auto child : childs)
            {
                std::vector<std::string> parents_of_child;
                for (auto& n : child->prev)
                    parents_of_child.push_back(n->name);
                if (contains_all(sorted_nodes, parents_of_child))
                {
                    sorted_nodes.push_back(child->name);
                    in_stock.push(child->name);
                }
            }
        }
        return sorted_nodes;
    }
}
